import torch
import torch.nn.functional as F
import torch.nn as nn
from .bconv import Bconv

class PDC(nn.Module):
    def __init__(self,channels, f4=True):
        super(PDC, self).__init__()
        self.f4 = f4

        # 采样倍数为2,进行两次采样
        self.upsameple = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)

        self.conv1 = Bconv(channels, channels, 3, padding=1)
        self.conv2 = Bconv(channels, channels, 3, padding=1)
        self.conv3 = Bconv(channels, channels, 3, padding=1)
        self.conv4 = Bconv(channels, channels, 3, padding=1)

        self.conv5 = Bconv(channels*2, channels*2, 3, padding=1)
        self.conv6 = Bconv(channels*2, channels*2, 3, padding=1)


        if f4:
            self.conv7 = Bconv(channels*4, channels*4, 3, padding=1)
            self.conv8 = Bconv(channels*4, channels*4, 3, padding=1)
            self.conv9 = nn.Conv2d(channels*4, 1, 1, padding=0)

        else:
            self.conv7 = Bconv(channels*3, channels*3, 3, padding=1)
            self.conv8 = Bconv(channels*3, channels*3, 3, padding=1)
            self.conv9 = nn.Conv2d(channels*3, 1, 1, padding=0)

    def forward(self, f1, f2, f3, f4=None):
        '''
        f1: B,C,W,H
        f2: B,C,W*2,H*2
        f3: B,C,W*4,H*4
        f4: B,C,W*4,H*4
        '''
        up_f1 = self.upsameple(f1)          # B,C,W*2,H*2
        up_f2 = self.upsameple(f2)          # B,C,W*4,H*4

        x11 = self.conv1(up_f1)* f2         # B,C,W*2,H*2
        x12 = self.conv2(up_f1)             # B,C,W*2,H*2
        x_cat1 = torch.cat((x11, x12), 1)   # B,C*2,W*2,H*2
        x_cat1 = self.conv5(x_cat1)         # B,C*2,W*2,H*2
        x_cat1 = self.upsameple(x_cat1)     # B,C*2,W*4,H*4
        x_cat1 = self.conv6(x_cat1)         # B,C*2,W*4,H*4

        x13 = self.upsameple(up_f1)         # B,C,W*4,H*4
        x13 = self.conv3(x13)               # B,C,W*4,H*4
        x21 = self.conv4(up_f2)             # B,C,W*4,H*4
        x_mul = x13*x21*f3                  # B,C,W*4,H*4

        if self.f4:
            x_cat2 = torch.cat((x_mul, x_cat1, f4), 1)  # B,C*4,W*4,H*4
            res = self.conv7(x_cat2)        # B,C*4,W*4,H*4
            res = self.conv8(res)           # B,C*4,W*4,H*4
            res = self.conv9(res)           # B,1,W*4,H*4
        else:
            x_cat2 = torch.cat((x_mul, x_cat1), 1)  # B,C*3,W*4,H*4
            res = self.conv7(x_cat2)        # B,C*3,W*4,H*4
            res = self.conv8(res)           # B,C*3,W*4,H*4
            res = self.conv9(res)           # B,1,W*4,H*4

        return res
    
if __name__ == '__main__':
    pdc1 = PDC(64)
    pdc2 = PDC(64, f4=False)
    f1 = torch.randn(1, 64, 32, 32)
    f2 = torch.randn(1, 64, 64, 64)
    f3 = torch.randn(1, 64, 128, 128)
    f4 = torch.randn(1, 64, 128,128)

    res1 = pdc1(f1, f2, f3, f4)
    res2 = pdc2(f1, f2, f3)

    print(res1.shape)
    print(res2.shape)